import os
import sys
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path

import fitz


def extract_images_from_pdf(pdf_path):
   """
   从单个PDF文件中提取图片到同名文件夹
   
   参数:
       pdf_path: PDF文件的路径
   """
   try:
       # 获取PDF文件名（不含扩展名）并创建同名文件夹
       pdf_path = Path(pdf_path)
       output_dir = pdf_path.parent / pdf_path.stem

       # 如果文件夹已存在，则跳过处理
       if output_dir.exists():
           print(f"文件夹已存在，跳过处理: {output_dir}")
           return

       output_dir.mkdir(exist_ok=True)

       # 打开PDF文件
       pdf_document = fitz.open(str(pdf_path))
       image_count = 0

       # 遍历每一页
       for page_num in range(pdf_document.page_count):
           page = pdf_document[page_num]
           images = page.get_images()

           # 遍历该页的所有图片
           for img_index, img in enumerate(images):
               try:
                   xref = img[0]
                   base_image = pdf_document.extract_image(xref)

                   if base_image:
                       image_bytes = base_image["image"]
                       image_ext = base_image["ext"]

                       # 生成图片文件名
                       image_name = f"{page_num + 1:03d}-{img_index + 1:03d}.{image_ext}"
                       image_path = output_dir / image_name

                       # 保存图片
                       with open(image_path, "wb") as image_file:
                           image_file.write(image_bytes)
                           image_count += 1

               except Exception as e:
                   print(
                       f"处理图片时出错 {pdf_path.name} - 页面 {page_num + 1}: {str(e)}")

       pdf_document.close()
       print(f"完成 {pdf_path.name} - 共提取 {image_count} 张图片")

   except Exception as e:
       print(f"处理PDF文件出错 {pdf_path}: {str(e)}")


def process_folder(folder_path):
   """
   处理文件夹中的所有PDF文件
   
   参数:
       folder_path: 包含PDF文件的文件夹路径
   """
   # 确保文件夹路径存在
   if not os.path.exists(folder_path):
       print(f"错误: 文件夹 '{folder_path}' 不存在")
       return

   # 获取所有PDF文件
   pdf_files = []
   for root, _, files in os.walk(folder_path):
       for file in files:
           if file.lower().endswith('.pdf'):
               pdf_files.append(os.path.join(root, file))

   if not pdf_files:
       print(f"在 {folder_path} 中未找到PDF文件")
       return

   print(f"找到 {len(pdf_files)} 个PDF文件，开始处理...")

   # 使用线程池并行处理PDF文件
   with ThreadPoolExecutor(max_workers=4) as executor:
       executor.map(extract_images_from_pdf, pdf_files)

   print("所有PDF处理完成！")


if __name__ == "__main__":
   # 指定要处理的文件夹路径
   folder_path = 'F:\\H图片\\美女图片\\清纯美女'  # 替换为你的文件夹路径

   # 如果通过命令行参数提供路径，则使用命令行参数
   if len(sys.argv) > 1:
       folder_path = sys.argv[1]

   process_folder(folder_path)
